package com.edu.learning.concurrent;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;

/**
 * @author ll
 * @Date 2017/11/1 16:48
 */
public class QueuableCachedThreadPool extends ThreadPoolExecutor {

    private final AtomicLong submittedCount = new AtomicLong(0);

    private final LongAdder longAdder = new LongAdder();

    public long getLongAdder() {
        return longAdder.longValue();
    }

    public long getSubmittedCount() {
        return submittedCount.get();
    }

    public QueuableCachedThreadPool(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
                                    ControllableQueue workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        workQueue.setParent(this);
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        submittedCount.decrementAndGet();
    }

    @Override
    public void execute(Runnable command) {
        execute(command, 0, TimeUnit.MILLISECONDS);
    }

    /**
     * Executes the given command at some time in the future. The command may execute in a new thread, in a pooled
     * thread, or in the calling thread, at the discretion of the <tt>Executor</tt> implementation. If no threads are
     * available, it will be added to the work queue. If the work queue is full, the system will wait for the specified
     * time and it throw a RejectedExecutionException if the queue is still full after that.
     *
     * TCC模式，预设资源分配成功，不成功补偿
     *
     * @param command the runnable task
     * @param timeout A timeout for the completion of the task
     * @param unit The timeout time unit
     * @throws RejectedExecutionException if this task cannot be accepted for execution - the queue is full
     * @throws NullPointerException if command or unit is null
     */
    public void execute(Runnable command, long timeout, TimeUnit unit) {
        submittedCount.incrementAndGet();
        try {
            super.execute(command);
        } catch (RejectedExecutionException rx) {
            final ControllableQueue queue = (ControllableQueue) super.getQueue();
            try {
                if (!queue.force(command, timeout, unit)) {
                    submittedCount.decrementAndGet();
                    throw new RejectedExecutionException("Queue capacity is full");
                }
            } catch (InterruptedException e) {
                submittedCount.decrementAndGet();
                throw new RejectedExecutionException(rx);
            }
        }
    }

    public static class ControllableQueue extends LinkedBlockingQueue<Runnable> {

        private static final long serialVersionUID = -2478296735394414999L;
        private volatile QueuableCachedThreadPool parent = null;

        public ControllableQueue(int capacity) {
            super(capacity);
        }

        public void setParent(QueuableCachedThreadPool queuableCachedThreadPool) {
            parent = queuableCachedThreadPool;
        }

        public boolean force(Runnable runnable) {
            if (parent.isShutdown()) {
                throw new RejectedExecutionException("Executor not running, can't force a command into the queue");
            }
            return super.offer(runnable);
        }

        public boolean force(Runnable runnable, long timeout, TimeUnit unit) throws InterruptedException {
            if (parent.isShutdown()) {
                throw new RejectedExecutionException("Executor not running, can't force a command into the queue");
            }
            return super.offer(runnable, timeout, unit);
        }

        @Override
        public boolean offer(Runnable runnable) {
            //threadPool.getPoolSize() 是个有锁的操作，所以尽量减少
            long currentPoolSize = parent.getPoolSize();

            if (currentPoolSize >= parent.getMaximumPoolSize()) {
                return super.offer(runnable);
            }
            if (parent.getSubmittedCount() < currentPoolSize) {
                return super.offer(runnable);
            }
            if (currentPoolSize < parent.getMaximumPoolSize()) {
                return false;
            }
            return super.offer(runnable);
        }
    }
}
